import gradio as gr
import numpy as np
import joblib
from PIL import Image
import os
import random
from util1 import extract_features  # Ensure util.py is in the same directory

class CatDogClassifier:
    def __init__(self, model_path, feature_method='vgg'):
        self.model = joblib.load(model_path)
        print(f"加载模型: {model_path},加载成功")
        self.feature_method = feature_method

    def preprocess_image(self, image):
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)

        if self.feature_method == 'flat':
            img_resized = image.resize((64, 64))
            img_array = np.array(img_resized)
            features = img_array.flatten().reshape(1, -1)
        else:  # 'vgg' 方法
            features = extract_features(image)
            if not np.any(features):
                print("警告: 提取的特征全为零")
            else:
                print("特征提取成功，包含非零值")
            print(f"提取的特征: {features}")

            # 确保特征为1024维
            if features.shape[1] == 512:
                features = np.concatenate([features, features], axis=1)
                print("特征已拼接为1024维")
            elif features.shape[1] < 512:
                padded_features = np.zeros((1, 512))
                padded_features[0, :features.shape[1]] = features
                features = np.concatenate([padded_features, padded_features], axis=1)
                print("特征已填充并拼接为1024维")
            elif features.shape[1] > 512:
                features = features[:, :512]
                features = np.concatenate([features, features], axis=1)
                print("特征已截断并拼接为1024维")

            print(f"调整后的特征形状: {features.shape}")
            print(f"调整后的特征示例: {features.flatten()[:10]}")

            # 归一化特征
            features = features.astype('float32')
            norm = np.linalg.norm(features)
            if norm != 0:
                features /= norm
            print(f"归一化后的特征示例: {features.flatten()[:10]}")

            # 进行预测
            try:
                prediction = self.model.predict(features)
                print(f"预测结果: {prediction}")
            except Exception as e:
                print(f"预测出错: {e}")

        return features



    def predict(self, image):
        features = self.preprocess_image(image)
        print(f"预测输入特征形状: {features.shape}")  # 添加输出
        probabilities = self.model.predict_proba(features)[0]
        print(f"预测概率: {probabilities}")  # 添加输出
        return {"猫": float(probabilities[0]), "狗": float(probabilities[1])}

# Create classifier instance
model_file = 'best_model_svm_20241031_223249.pkl'
classifier = CatDogClassifier(model_path=model_file, feature_method="vgg")

# Define prediction function
def predict_image(image):
    if image is None:
        return None
    try:
        return classifier.predict(image)
    except Exception as e:
        print(f"预测出错: {str(e)}")
        return None

# Create Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="上传图片"),
    #outputs=gr.Label(num_top_classes=2, label="预测结果"),
    outputs=gr.Label(label="预测结果"),

    title="🐱猫狗图片分类器🐶",
    description="""\
    ## 使用说明
    1. 点击上传或拖拽一张包含猫或狗的图片
    2. 等待AI预测结果
    3. 查看预测结果和置信度

    *支持的图片格式：JPG、PNG、JPEG*
    """,
    examples=None,  # Remove random examples since they are not needed
    cache_examples=True
)

# Start the application
if __name__ == "__main__":
    print(f"使用模型: {model_file}")
    iface.launch(
        server_name="0.0.0.0",  # Allow external access
        share=True,              # Create a public link
        server_port=7860,       # Specify port number
        debug=True               # Debug mode
    )
